import tensorflow as tf
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
import sys
import cv2
from PIL import Image
import numpy as np


class MainWindow(QTabWidget):
    def __init__(self):
        super().__init__()
        self.setWindowIcon(QIcon('images/1.png'))
        self.setWindowTitle('花卉识别系统')

        # 使用cnn模型
        self.model = tf.keras.models.load_model("models/cnn_flower.h5")
        self.to_predict_name = "images/init.jpg"
        # 花的种类
        self.class_names = ['风铃草','毛茛','款冬','牛至','藏红花','水仙花','雏菊','蒲公英','贝母','鸢尾花','百合','三色堇','雪花莲','向日葵','虎尾兰','郁金香','风信子']
        self.class_enames =['Bluebell','Buttercup', "Colls' Foot", 'Cowslip','Crocus', 'Daffodil','Daisy', 'Dandelion', 'Fritillary ', 'Iris', 'Lily Valley', 'Pansy', 'Snowdrop', 'Sunflower', 'Tigerlily', 'Tulip', 'Windflower']
        self.resize(1000, 900)
        self.initUI()

    def initUI(self):
        main_widget = QWidget()
        main_layout = QHBoxLayout()
        font = QFont('宋体', 15)

        left_widget = QWidget()
        left_layout = QVBoxLayout()
        self.img_title = QLabel("测试样本")
        self.img_title.setFont(QFont('宋体', 16))

       # img_title.setAlignment(Qt.AlignCenter)
        self.img_label = QLabel()
        # img_init = cv2.imread(self.to_predict_name)
        img_init = cv2.imdecode(np.fromfile(self.to_predict_name, dtype=np.uint8), -1)
        img_init = cv2.resize(img_init, (224, 224))
        cv2.imwrite('images/target.png', img_init)
        self.img_label.setPixmap(QPixmap('images/target.png'))
        #left_layout.addWidget(self.img_label)
        left_layout.addStretch()

        left_layout.addWidget(self.img_label, 0, Qt.AlignCenter)
        left_layout.addWidget(self.img_title, 0, Qt.AlignCenter)
        left_layout.addStretch()
        # left_layout.setAlignment(Qt.AlignCenter)
        # left_widget.setLayout(left_layout)

        right_widget = QWidget()
        right_layout = QVBoxLayout()
        btn_change = QPushButton(" 选择图片 ")
        btn_change.clicked.connect(self.change_img)
        btn_change.setFont(font)

        btn_predict = QPushButton(" 识 别  ")
        btn_predict.setFont(font)
        btn_predict.clicked.connect(self.predict_img)

        label_result = QLabel(' 识 别 结 果 ')
        self.result = QLabel("待识别")
        label_result.setFont(QFont('宋体', 16))
        self.result.setFont(QFont('宋体', 24))

        self.eresult = QLabel("")
        self.eresult.setFont(QFont('宋体', 24))

        right_layout.addStretch()
        right_layout.addWidget(label_result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addWidget(self.eresult, 0, Qt.AlignCenter)
        #right_layout.addStretch()
        right_layout.addWidget(self.result, 0, Qt.AlignCenter)


        right_layout.addStretch()

        # 加载测试图按钮，识别按钮：
        right_layout.addWidget(btn_change)
        right_layout.addWidget(btn_predict)

        right_layout.addStretch()
        # right_layout.addSpacing(5)
        # right_widget.setLayout(right_layout)
        right_widget.setLayout(left_layout)
        left_widget.setLayout(right_layout)
        # 关于页面
        about_widget = QWidget()
        about_layout = QVBoxLayout()
        about_title = QLabel('花卉模型训练结果')
        about_title.setFont(QFont('宋体', 18))
        about_title.setAlignment(Qt.AlignCenter)
        about_img = QLabel()


        # 热力图显示
        about_img.setPixmap(QPixmap('./results/heatmap_cnn.png'))

        about_img.setAlignment(Qt.AlignCenter)
        label_super = QLabel()

        label_super.setFont(QFont('宋体', 12))
        label_super.setOpenExternalLinks(True)
        label_super.setAlignment(Qt.AlignRight)
        # git_img = QMovie('images/')
        about_layout.addWidget(about_title)
        about_layout.addStretch()
        about_layout.addWidget(about_img)

        about_layout.addStretch()
        about_layout.addWidget(label_super)

        # 跳转到训练结果
        about_widget.setLayout(about_layout)

        main_layout.addWidget(left_widget)
        main_layout.addWidget(right_widget)

        # 跳转到主界面，默认
        main_widget.setLayout(main_layout)
        self.addTab(main_widget, '主页面')
        self.addTab(about_widget, '训练结果')
        self.setTabIcon(0, QIcon('images/主页面.png'))
        self.setTabIcon(1, QIcon('images/关于.png'))
    # 选择图片相应事件
    def change_img(self):
        try:
            openfile_name = QFileDialog.getOpenFileName(self, '选择文件', '', 'Image files(*.jpg , *.png)')
            # print(openfile_name)

            img_name = openfile_name[0]
            if img_name == '':
                pass
            else:
                self.to_predict_name = img_name
                # img_init = cv2.imread(self.to_predict_name)
                img_init = cv2.imdecode(np.fromfile(self.to_predict_name, dtype=np.uint8), -1)
                # 图片像素更改为前面模型统一的224*224，because模型接收的参数为224*224
                img_init = cv2.resize(img_init, (224, 224))
                cv2.imwrite('images/target.png', img_init)
                self.img_label.setPixmap(QPixmap('images/target.png'))
                self.result.setText("待识别")
                self.eresult.setText("")
        except Exception as e:
            print("选择测试样本异常:" + repr(e))

    # 识别相应事件
    def predict_img(self):
        try:
            img = Image.open('images/target.png')
            img = np.asarray(img)
            # gray_img = img.convert('L')
            # img_torch = self.transform(gray_img)
            outputs = self.model.predict(img.reshape(1, 224, 224, 3))
            # 224*224的图片数组转换为np数组，得到索引
            result_index = np.argmax(outputs)
            # print(result_index)
            eresult = self.class_enames[result_index]
            result = self.class_names[result_index]
            self.eresult.setText(eresult)
            self.result.setText(result)

            print("识别结果:"+str(eresult)+'=>'+str(result))
        except Exception as e:
            print("花卉识别异常:" + repr(e))

    def closeEvent(self, event):
        """
        关闭确认框提示
        :param event:
        """
        quit_flag = QMessageBox.question(self, '退出确认', '确认退出应用程序?', QMessageBox.Yes | QMessageBox.No)
        event.ignore()
        if quit_flag == QMessageBox.Yes:
            event.accept()
            print("已经退出")
            super().closeEvent(event)
            sys.exit()


if __name__ == "__main__":
    try:
        app = QApplication(sys.argv)
        x = MainWindow()
        x.show()
        sys.exit(app.exec_())

    except Exception as e:
        print("异常信息:" + repr(e))
        sys.exit(app.exec_())
